/*
 * GraphState2.cpp
 *
 *  Created on: Jun 28, 2011
 *      Author: pickrell
 */

#include "TreeState.h"

TreeState::TreeState(){
}

TreeState::TreeState(Alignment * al, PhyloPop_params *p){
	params = p;
	aln = al;
	allhapnames = al->list_haps();
	unsigned int seed = unsigned( time(NULL));
	//seed = 1348696974;
	cout << "SEED: "<< seed << "\n";
	srand ( seed );
	random_shuffle(allhapnames.begin(), allhapnames.end() );
	vector<string> startpops;
	startpops.push_back(allhapnames[0]); startpops.push_back(allhapnames[1]); startpops.push_back(allhapnames[2]);
	vector<vector<bool> > starthaps;
	vector<int> colcounts;
	for (int i = 0; i < aln->nsite; i++) colcounts.push_back(aln->colcounts.at(i));
	starthaps.push_back(aln->get_hap(startpops[0])); starthaps.push_back(aln->get_hap(startpops[1])); starthaps.push_back(aln->get_hap(startpops[2]));
	tree = new PopTree(startpops, starthaps, aln->len, aln->nmono, colcounts, params);
	tree_bk = new PopTree(startpops, starthaps, aln->len,aln->nmono, colcounts, params);
	tree_bk2 = new PopTree(startpops, starthaps, aln->len, aln->nmono, colcounts,params);
	tree_bk3 = new PopTree(startpops, starthaps, aln->len, aln->nmono, colcounts,params);
	current_nhap = 3;
	tree->compute_tree_llk();
	current_llik = tree->current_llk;
	//vector<Graph::vertex_descriptor> inorder = tree->get_inorder_traversal(3);

	local_hillclimb(3);
	cout << "Starting from:\n";
	cout << tree->get_newick_format() << "\n";
	cout << "llk: "<< tree->current_llk << "\n";
}


void TreeState::add_pop(){

	string toadd = allhapnames[current_nhap];
	cout << "Adding "<< toadd << "\n";
	vector<bool> haptoadd = aln->get_hap(toadd);
	vector<Graph::vertex_descriptor> postorder = tree->get_postorder_traversal(current_nhap);
	int max_index;
	double max_llik;
	tree_bk->copy(tree);

	for (int i = 0; i < postorder.size(); i++){

		Graph::vertex_descriptor tmp = tree->add_tip(postorder[i], haptoadd, toadd);
		current_nhap++;

		tree->compute_tree_llk(tmp);
		//cout << i << " "<< tree->g[postorder[i]].index <<" "<< tree->current_llk << "\n";
		double llk = tree->current_llk;
		if (i == 0){
			max_index = i;
			max_llik = llk;
		}
		else if (llk > max_llik){
			max_index = i;
			max_llik = llk;
		}
		tree->copy(tree_bk);
		current_nhap--;
		postorder = tree->get_postorder_traversal(current_nhap);
	}
	Graph::vertex_descriptor tmp = tree->add_tip(postorder[max_index], haptoadd, toadd);
	current_nhap++;
	tree->compute_tree_llk(tmp);
	current_llik = tree->current_llk;

}

void TreeState::iterate_hillclimb(){
	//cout << "here\n"; cout.flush();
	int changes = many_local_hillclimb();
	while (changes > 0) 	{
		//cout << changes << " changes\n"; cout.flush();
		changes = many_local_hillclimb();
	}
}

int TreeState::many_local_hillclimb(){
	int leninorder = 2*current_nhap -1;
	int toreturn = 0;
	for (int i = 1; i < leninorder; i++){
		//cout << i << "\n";
		toreturn += local_hillclimb(i);
	}
	return toreturn;
}

int TreeState::local_hillclimb(int postorder_index){
	// if there was a rearrangement, return 1. otw 0.
	//
	vector<Graph::vertex_descriptor> postorder = tree->get_postorder_traversal(current_nhap);
	if ( tree->g[ postorder[postorder_index]].is_root) return local_hillclimb_root();
	if (tree->g[postorder[postorder_index]].is_tip) return 0;
	double llik1, llik2, llik3;
	//cout << "here1\n"; cout.flush();
	tree_bk->copy(tree);


	llik1 = current_llik;

	tree->local_rearrange(postorder[postorder_index], 1);
	//cout << "here2\n"; cout.flush();
	tree->compute_tree_llk();
	//cout << "here3\n"; cout.flush();
	llik2 =  tree->current_llk;

	tree_bk2->copy(tree);


	tree->copy(tree_bk);
	postorder = tree->get_postorder_traversal(current_nhap);
	tree->local_rearrange(postorder[postorder_index], 2);

	tree->compute_tree_llk();
	llik3 =  tree->current_llk;

	tree_bk3->copy(tree);
	tree->copy(tree_bk);

	//cout << postorder_index << " "<< llik1 <<" "<< llik2 << " "<< llik3 << "\n";
	if (llik1 < llik2 || llik1 < llik3){
		if (llik2 > llik3){
			tree->copy(tree_bk2);
			current_llik = llik2;
			return 1;
		}
		else{
			tree->copy(tree_bk3);
			current_llik = llik3;
			return 1;
		}
	}
	return 0;
}

void TreeState::set_to_newick(string infile){
	igzstream in(infile.c_str());
	string st;
	getline(in, st);
	vector<string> line;
	string buf;
	stringstream ss(st);
	line.clear();
	while (ss>> buf){
		line.push_back(buf);
	}
	string newick = line[0];
	tree->set_graph(newick);
	current_nhap = tree->ntip;
	set_coaltimes();
}

int TreeState::local_hillclimb_root(){
	double best_llik = current_llik;
	tree_bk->copy(tree);
	tree_bk2->copy(tree);

	int toreturn = 0;
	for (int i = 1; i <= 4; i++){
		tree->move_root(i);
		tree->compute_tree_llk();

		double tmplik = tree->current_llk;
		if (tmplik > best_llik){
			toreturn = 1;
			tree_bk->copy(tree);
			best_llik = tmplik;
		}
		tree->copy(tree_bk2);
	}
	if (toreturn == 1){
		tree->copy(tree_bk);
		current_llik = best_llik;
	}
	return toreturn;

}


void TreeState::print_coales(){
	set_labeled_coaltimes();
	for(vector<double>::iterator it = coaltimes1.begin(); it != coaltimes1.end(); it++){
		cout << *it << " 1\n";
	}
	for(vector<double>::iterator it = coaltimes2.begin(); it != coaltimes2.end(); it++){
		cout << *it << " 2\n";
	}
	for(vector<double>::iterator it = coaltimes3.begin(); it != coaltimes3.end(); it++){
		cout << *it << " 3\n";
	}

}

void TreeState::print_coals12(){
	set_labeled_coaltimes();
	for (vector<pair<double, int> >::iterator it = coals12.begin(); it != coals12.end(); it++){
		cout << it->first << " "<< it->second << "\n";
	}
}


void TreeState::set_coaltimes(){
	coaltimes.clear();
	vector<Graph::vertex_descriptor> v= tree->get_inorder_traversal(current_nhap);
	//cout << "here2\n"; cout.flush();
	for (vector<Graph::vertex_descriptor>::iterator it = v.begin(); it != v.end(); it++){
		if (tree->g[*it].is_tip == false){
			float tmp = tree->g[*it].height;
			coaltimes.push_back(tmp);
		}
	}
	sort(coaltimes.begin(), coaltimes.end());
}

void TreeState::set_labeled_coaltimes(){
	coaltimes1.clear();
	coaltimes2.clear();
	coaltimes3.clear();
	coals12.clear();
	n1 = 0;
	n2 = 0;
	vector<Graph::vertex_descriptor> v= tree->get_inorder_traversal(current_nhap);
	for (vector<Graph::vertex_descriptor>::iterator it = v.begin(); it != v.end(); it++){
		if (tree->g[*it].is_tip == false){
			if (tree->g[*it].pop == 1) {
				coaltimes1.push_back(tree->g[*it].height);
				coals12.push_back(make_pair(tree->g[*it].height, tree->g[*it].pop));
				//coals12.push_back(*it);
			}
			else if (tree->g[*it].pop == 2) {
				coals12.push_back(make_pair(tree->g[*it].height, tree->g[*it].pop));
				coaltimes2.push_back(tree->g[*it].height);
			}
			else if (tree->g[*it].pop == 3) coaltimes3.push_back(tree->g[*it].height);
		}
		else{
			if (tree->g[*it].pop  == 1) n1++;
			else if (tree->g[*it].pop == 2) n2++;
		}
	}
	sort(coaltimes1.begin(), coaltimes1.end());
	sort(coaltimes2.begin(), coaltimes2.end());
	sort(coaltimes3.begin(), coaltimes3.end());
	sort(coals12.begin(), coals12.end(), comp_vertex);
}

bool comp_vertex(pair<double, int> v1, pair<double, int> v2){
	return( v1.first < v2.first);
}

void TreeState::init_tip_labels(double h, gsl_rng * r){
	/*
	 * initialize labels at random with a given split time
	 *
	 */
	Graph::vertex_descriptor best;
	double besth;
	vector<Graph::vertex_descriptor> inorder = tree->get_inorder_traversal(current_nhap);
	best = inorder[0];
	besth = 0;
	for (vector<Graph::vertex_descriptor>::iterator it = inorder.begin(); it != inorder.end(); it++){
		double testh = tree->g[*it].height;
		//if (tree->g[*it].index == 10) cout << "10 "<< tree->g[*it].height << "\n";
		//if (tree->g[*it].index == 0) cout << "0 "<< tree->g[*it].height << "\n";
		if (besth < h && testh > besth){
			besth = testh;
			best = *it;
		}
		else if (besth > h && testh < besth && testh > h){
			besth = testh;
			best = *it;
		}
	}
	if (besth < h){
		//cout << "here\n";
		//cout << tree->g[best].index << "\n";
		//cout << tree->g[best].is_root << "\n";
		double ran = gsl_rng_uniform(r);
		if (ran < 0.5) {
			assign_labels(best, 1);
			while(!tree->g[best].is_root){

				best = tree->get_parent_node(best).first;
				assign_labels(best, 1);
			}
		}
		else {
			assign_labels(best, 2);
			while(!tree->g[best].is_root){

				best = tree->get_parent_node(best).first;
				assign_labels(best, 2);
			}
		}
	}
	else{
		pair<Graph::vertex_descriptor, Graph::vertex_descriptor> ch = tree->get_child_nodes(best);
		double ran = gsl_rng_uniform(r);
		if (ran < 0.5){
			assign_labels(ch.first, 1);
			assign_labels(ch.second, 2);
		}
		else{
			assign_labels(ch.first, 2);
			assign_labels(ch.second, 1);

		}
		while (!tree->g[best].is_root){
			assign_labels(best, r, h);
			best = tree->get_parent_node(best).first;
		}
		//cout << tree->g[best].index << "\n";
		assign_labels(best, r, h);
	}
	set_labeled_coaltimes();
}

void TreeState::assign_labels(Graph::vertex_descriptor v, int label){
	tree->g[v].pop = label;
	if (!tree->g[v].is_tip){
		pair<Graph::vertex_descriptor, Graph::vertex_descriptor> ch = tree->get_child_nodes(v);
		assign_labels(ch.first, label);
		assign_labels(ch.second, label);
	}
}

void TreeState::flip_labels(vector<Graph::vertex_descriptor> toflip){
	for (vector<Graph::vertex_descriptor>::iterator it = toflip.begin(); it != toflip.end(); it++){
		if (tree->g[*it].pop == 1) tree->g[*it].pop = 2;
		else if (tree->g[*it].pop == 2) tree->g[*it].pop = 1;
	}
	set_labeled_coaltimes();
}

vector<Graph::vertex_descriptor> TreeState::move_label_subset(vector<Graph::vertex_descriptor> toflip, int which){
	vector<Graph::vertex_descriptor> toreturn;
	for (vector<Graph::vertex_descriptor>::iterator it = toflip.begin(); it != toflip.end(); it++){
		if (tree->g[*it].pop == which and which ==1) {
			toreturn.push_back(*it);
			tree->g[*it].pop = 2;
		}
		else if (tree->g[*it].pop == which and which ==2) {
			toreturn.push_back(*it);
			tree->g[*it].pop = 1;
		}
	}
	set_labeled_coaltimes();
	return toreturn;
}

void TreeState::assign_labels(Graph::vertex_descriptor v, gsl_rng* r, double h){
	if (tree->g[v].height > h ){
		tree->g[v].pop = 3;
		pair<Graph::vertex_descriptor, Graph::vertex_descriptor> ch = tree->get_child_nodes(v);
		//cout << tree->g[ch.first].index << " "<< tree->g[ch.first].popid << " "<< tree->g[ch.second].index << " "<< tree->g[ch.second].popid << "\n";
		if (tree->g[ch.first].pop < 1) {
			//cout << tree->g[v].index << " "<< tree->g[ch.first].index << "\n";
			assign_labels(ch.first, r, h);
		}
		if (tree->g[ch.second].pop < 1) {
			//cout << tree->g[v].index << " "<< tree->g[ch.first].index << "\n";
			assign_labels(ch.second, r, h);
		}
	}
	else{
		double ran = gsl_rng_uniform(r);
		if (ran < 0.5) assign_labels(v, 1);
		else assign_labels(v, 2);
	}
}

void TreeState::optimize_all_tips(double p1, vector<double> * bpoints12, vector<double> * popsizes1, vector<double> *popsizes2){
	vector<Graph::vertex_descriptor> postorder = tree->get_postorder_traversal(tree->ntip);
	for (vector<Graph::vertex_descriptor>::iterator it = postorder.begin(); it != postorder.end(); it++){
		if (tree->g[*it].pop == 3) optimize_single_tip(*it, p1, bpoints12, popsizes1, popsizes2);
	}
}

void TreeState::iterate_optimize_all_tips(double p1, vector<double> * bpoints12, vector<double> * popsizes1, vector<double> *popsizes2){
	double st_llk = structured_llk_no3(p1, bpoints12, popsizes1, popsizes2);
	//cout << "start "<< st_llk << "\n";
	bool go = true;
	while (go){
		//cout << "here\n"; cout.flush();
		optimize_all_tips(p1, bpoints12, popsizes1, popsizes2);
		double testllk = structured_llk_no3(p1, bpoints12, popsizes1, popsizes2);
		//cout << "test "<< testllk << "\n";
		if (testllk < st_llk +params->e_lk) go = false;
		else st_llk = testllk;
	}
}

void TreeState::optimize_single_tip(Graph::vertex_descriptor v, double p1, vector<double> * bpoints12, vector<double> * popsizes1, vector<double> *popsizes2){
	if (tree->g[v].pop != 3){
		cerr<< "ERROR: trying to optimize tip "<< tree->g[v].index << " with popid "<< tree->g[v].pop<< "\n";
		exit(1);
	}
	pair<Graph::vertex_descriptor, Graph::vertex_descriptor> ch = tree->get_child_nodes(v);
	map<string, Graph::vertex_descriptor> tips_left = tree->get_tips(ch.first);
	map<string, Graph::vertex_descriptor> tips_right = tree->get_tips(ch.second);
	double st_llk = structured_llk_no3(p1, bpoints12, popsizes1, popsizes2);
	double bestllk = st_llk;
	int best = 0;
	bool hasboth = false;
	bool has1 = false;
	bool has2 = false;
	bool left_hasboth = false;
	bool left_has1 = false;
	bool left_has2 = false;
	bool right_hasboth = false;
	bool right_has1 = false;
	bool right_has2 = false;

	int numtip = 0;
	// testing whether the left side and the right side have 1s and 2s
	for (map<string, Graph::vertex_descriptor>::iterator it = tips_left.begin(); it != tips_left.end(); it++){
		if( tree->g[it->second].pop == 1) {
			has1 = true;
			left_has1 = true;
		}
		if( tree->g[it->second].pop == 2) {
			left_has2 = true;
			has2 = true;
		}
		numtip++;
	}
	for (map<string, Graph::vertex_descriptor>::iterator it = tips_right.begin(); it != tips_right.end(); it++){
		if( tree->g[it->second].pop == 1) {
			right_has1 = true;
			has1 = true;
		}
		if( tree->g[it->second].pop == 2) {
			right_has2 = true;
			has2 = true;
		}
		numtip++;
	}
	if (has1 && has2) hasboth =true;
	//cout << "start "<< st_llk << "\n";

	vector<Graph::vertex_descriptor> toflip = tree->get_postorder_traversal(v, numtip);
	vector<Graph::vertex_descriptor> flip_left = tree->get_postorder_traversal(ch.first, tips_left.size());
	vector<Graph::vertex_descriptor> flip_right = tree->get_postorder_traversal(ch.second, tips_right.size());

	if (left_has1){
		vector<Graph::vertex_descriptor> moved = move_label_subset(flip_left, 1);
		double lk1 = structured_llk_no3(p1, bpoints12, popsizes1, popsizes2);
		//cout << "trying left1 "<< lk1 << "\n"; cout.flush();
		//tree->print("left1");
		if (lk1 > bestllk ){
			best = 1;
			bestllk = lk1;
		}
		move_label_subset(moved, 2);
	}
	if (left_has2){
		vector<Graph::vertex_descriptor> moved = move_label_subset(flip_left, 2);
		double lk2 = structured_llk_no3(p1, bpoints12, popsizes1, popsizes2);
		//cout << "trying left2 "<< lk2 << "\n"; cout.flush();
		//tree->print("left2");
		if (lk2 > bestllk ){
			best = 2;
			bestllk = lk2;
		}
		move_label_subset(moved, 1);
	}

	if (right_has1){
		vector<Graph::vertex_descriptor> moved = move_label_subset(flip_right, 1);
		double lk3 = structured_llk_no3(p1, bpoints12, popsizes1, popsizes2);
		//cout << "trying right1 "<< lk3 << "\n"; cout.flush();
		//tree->print("right1");
		if (lk3 > bestllk ){
			best = 3;
			bestllk = lk3;
		}
		move_label_subset(moved, 2);
	}
	if (right_has2){
		vector<Graph::vertex_descriptor> moved = move_label_subset(flip_right, 2);
		double lk4 = structured_llk_no3(p1, bpoints12, popsizes1, popsizes2);
		//cout << "trying right2 "<< lk4 << "\n"; cout.flush();
		//tree->print("right2");
		if (lk4 > bestllk ){
			best = 4;
			bestllk = lk4;
		}
		move_label_subset(moved, 1);
	}

	if (hasboth){
		flip_labels(toflip);
		double lk5 = structured_llk_no3(p1, bpoints12,  popsizes1, popsizes2);
		if (lk5 > bestllk ){
			best = 5;
			bestllk = lk5;
		}

		flip_labels(toflip);
	}
	if (best ==1) move_label_subset(flip_left, 1);
	else if (best ==2) move_label_subset(flip_left, 2);
	else if (best ==3) move_label_subset(flip_right, 1);
	else if (best ==4) move_label_subset(flip_right, 2);
	else if (best ==5) flip_labels(toflip);
	//cout << "after "<< n1 << " "<< n2 << " "<< coaltimes1.size() << " "<< coaltimes2.size() << "\n";
}
void TreeState::set_known(int cutoff){
	map<string, Graph::vertex_descriptor> tips = tree->get_tips(tree->root);
	for (map<string, Graph::vertex_descriptor>::iterator it = tips.begin(); it != tips.end(); it++){
		string id = tree->g[it->second].name;
		int popid = atoi(id.c_str());
		if (popid < cutoff){
			set_popid(it->second, 1);
		}
		else set_popid(it->second, 2);
	}
	set_labeled_coaltimes();
}

void TreeState::set_popid(Graph::vertex_descriptor v, int which){
	if (tree->g[v].pop != which){
		tree->g[v].pop = which;
		Graph::vertex_descriptor p = tree->get_parent_node(v).first;
		while (tree->g[p].pop != 3 && !tree->g[p].is_root){
			tree->g[p].pop = which;
			p = tree->get_parent_node(p).first;
		}
	}
}

double TreeState::llk(vector<double>* bpoints, vector<double>* popsizes){
	double toreturn = 0;
	int i = 0; //index coaltimes
	int j = 0; //index bpoints
	int k = 0; //index popsizes
	double previous_coaltime = 0;
	double previous_bpoint = 0;
	double previous_popsize = popsizes->at(0);
	while (i < coaltimes.size()){
	//while (i < 10){
		double current_coaltime = coaltimes.at(i);
		double current_bpoint = bpoints->at(j);
		double current_popsize = popsizes->at(k);
		int nlin = coaltimes.size()-i+1;
		//cout << i << " "<< nlin <<" "<<  current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";
		double add = 0;
		double logC = gsl_sf_lnfact(nlin) - gsl_sf_lnfact(2) - gsl_sf_lnfact(nlin-2);
		//cout << i << " C "<< C <<" nlin: "<< nlin << " "<< coaltimes.size()<< "\n";
		while (current_bpoint < current_coaltime){

			//add the bit after the previous event
			if (previous_coaltime < current_bpoint && previous_coaltime > previous_bpoint)	{
					add += - (current_bpoint - previous_coaltime) * exp(logC) / current_popsize;

			}
			else	{
				add += - (current_bpoint - previous_bpoint)* exp(logC) / current_popsize;
			}
			//increment
			j++;
			k++;
			previous_bpoint = current_bpoint;
			previous_popsize = current_popsize;
			current_bpoint = bpoints->at(j);
			current_popsize = popsizes->at(k);

		}
		if (previous_coaltime > previous_bpoint){
			//cout << "here4 "<< current_coaltime << " "<< previous_coaltime << " "<< C << " "<< current_popsize << " \n";
			//cout << i  << " adding5 " <<-(current_coaltime - previous_coaltime)*  C/ current_popsize << "\n";
			add += -(current_coaltime - previous_coaltime)*  exp(logC)/ current_popsize;
		}
		else {
			//cout << i << " adding6 " << - (current_coaltime - previous_bpoint) * C /current_popsize << "\n";
			add += - (current_coaltime - previous_bpoint) * exp(logC) /current_popsize;
		}
		//cout << i << " "<< add << "\n";
		add += logC - log(current_popsize);
		//cout << "add "<<i << " "<< add << "\n";
		toreturn += add;
		//cout << i << " "<< add << "\n";
		previous_coaltime = current_coaltime; //increment, save coalescent time
		i = i+1;
	}
	return toreturn;

}

double TreeState::structured_llk(double mergetime, double f1, vector<double> * bpoints12,  vector<double> *bpoints3, vector<double> *popsizes1, vector<double> *popsizes2, vector<double> *popsizes3){
	double toreturn = 0;
	toreturn += structured_llk_binom(f1);
	toreturn += llk12(bpoints12, popsizes1, popsizes2);
	//toreturn += llk(n1, 0, 0, coaltimes1, bpoints1, popsizes1);
	//toreturn += llk(n2, 0, 0, coaltimes2, bpoints2, popsizes2);
	toreturn += llk(coaltimes3.size()+1, mergetime, mergetime, coaltimes3, bpoints3, popsizes3 );
	return toreturn;
}

double TreeState::structured_llk_nof(double mergetime, vector<double> * bpoints12,  vector<double> *bpoints3, vector<double> *popsizes1, vector<double> *popsizes2, vector<double> *popsizes3){
	double toreturn = 0;
	toreturn += llk12(bpoints12, popsizes1, popsizes2);
	//toreturn += llk(n1, 0, 0, coaltimes1, bpoints1, popsizes1);
	//toreturn += llk(n2, 0, 0, coaltimes2, bpoints2, popsizes2);
	toreturn += llk(coaltimes3.size()+1, mergetime, mergetime, coaltimes3, bpoints3, popsizes3 );
	return toreturn;
}

double TreeState::structured_llk_no3(double f1, vector<double> * bpoints, vector<double> *popsizes1, vector<double> *popsizes2){
	double toreturn = 0;
	//cout << n1 << " "<< n2 << " "<< f1 << " "<< structured_llk_binom(f1) << "\n";
	toreturn += structured_llk_binom(f1);
	//cout << "coalpart "<< llk12(bpoints, popsizes1, popsizes2) << "\n";
	toreturn += llk12(bpoints, popsizes1, popsizes2);
	//toreturn += llk(n1, 0, 0, coaltimes1, bpoints, popsizes1);
	//toreturn += llk(n2, 0, 0, coaltimes2, bpoints, popsizes2);
	//cout << toreturn << "\n";
	return toreturn;
}

double TreeState::llk(int st_nlin, double st_coaltime, double st_bpoint, vector<double> coaltimes_in, vector<double>* bpoints, vector<double>* popsizes){

	// From Griffiths and Tavare (1994). Good summary in Heled and Drummond (2008)
	//
	// l(T| N(t) = \prod_{i=2}^n \frac{i choose 2}{N(t_i)} \exp {- \int_{t_{i+1}}^t_i \frac{i choose 2}{N(t)} dt}


	double toreturn = 0;
	int i = 0; //index coaltimes
	int j = 0; //index bpoints
	int k = 0; //index popsizes
	double previous_coaltime = st_coaltime;
	double previous_bpoint = st_bpoint;
	double previous_popsize = popsizes->at(0);
	while (i < coaltimes_in.size()){
		//cout << i << " "<< coaltimes_in.size() << " "<< j << " "<<bpoints->size() << " "<< k << " "<< popsizes->size()<< "\n"; cout.flush();
		double current_coaltime = coaltimes_in.at(i);
		double current_bpoint = bpoints->at(j);
		double current_popsize = popsizes->at(k);
		int nlin = st_nlin-i;
		//cout << i << " "<< current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";
		double add = 0;
		//cout << nlin<< " "<< st_nlin << " "<< coaltimes_in.size() << " "<< i <<"\n";
		//tree->print("test");
		double C = gsl_sf_fact(nlin) / (gsl_sf_fact(2)* gsl_sf_fact(nlin-2));
		//cout << i << " C "<< C <<" nlin: "<< nlin << "\n";
		while (current_bpoint < current_coaltime){

			//add the bit after the previous event
			if (previous_coaltime < current_bpoint && previous_coaltime > previous_bpoint)	{

					add += - (current_bpoint - previous_coaltime) * C / current_popsize;

			}
			else	{
				add += - (current_bpoint - previous_bpoint)* C / current_popsize;
			}
			//increment
			j++;
			k++;
			previous_bpoint = current_bpoint;
			previous_popsize = current_popsize;
			current_bpoint = bpoints->at(j);
			current_popsize = popsizes->at(k);

		}
		if (previous_coaltime > previous_bpoint){
			//cout << "here4 "<< current_coaltime << " "<< previous_coaltime << " "<< C << " "<< current_popsize << " \n";
			//cout << i  << " adding5 " <<-(current_coaltime - previous_coaltime)*  C/ current_popsize << "\n";
			add += -(current_coaltime - previous_coaltime)*  C/ current_popsize;
		}
		else {
			//cout << i << " adding6 " << - (current_coaltime - previous_bpoint) * C /current_popsize << "\n";
			add += - (current_coaltime - previous_bpoint) * C /current_popsize;
		}
		//cout << i << " adding7 " << log(C) - log(current_popsize) << "\n";
		add += log(C) - log(current_popsize);
		toreturn += add;
		//cout << i << " "<< add << "\n";
		previous_coaltime = current_coaltime; //increment, save coalescent time
		i = i+1;
	}
	// now all modeled coalescents have happened
	// if there is still more than one lineage, need to add likelihood of *no* additional coalescences in the remaining time
	//
;
	if (st_nlin > coaltimes_in.size()+1){

		double add = 0;
		double lastcoal;
		if (coaltimes_in.size() ==0 ) lastcoal = st_coaltime;
		else lastcoal= coaltimes_in.at(coaltimes_in.size()-1);
		int nlin = st_nlin -coaltimes_in.size();
		double C = gsl_sf_fact(nlin) / (gsl_sf_fact(2)* gsl_sf_fact(nlin-2));
		for (int i = 0; i < bpoints->size(); i++){
			double dt;
			double N = popsizes->at(i);
			if (bpoints->at(i) < lastcoal) continue;
			if (i > 0 && bpoints->at(i) > lastcoal && bpoints->at(i-1) < lastcoal){
				dt = bpoints->at(i) - lastcoal;
			}
			else if (i > 0){
				dt = bpoints->at(i) - bpoints->at(i-1);
			}
			add +=  - dt * C/ N;
		}
		toreturn += add;
	}

	return toreturn;

}


double TreeState::llk12( vector<double>* bpoints, vector<double>* popsizes1, vector<double>* popsizes2){

	// Structured likelihood with variable population sizes.
	// Single populatione equivalents are in Griffiths and Tavare (1994). Good summary in Heled and Drummond (2008)
	//
	// in coals12 are the coalescents times and labels
	// l( T, L | N1(t), N2(t)) = \prod{i = 2}^n lk(i)
	//


	double toreturn = 0;
	int i = 0; //index coaltimes
	int j = 0; //index bpoints
	int k = 0; //index popsizes
	double previous_coaltime = 0;
	double previous_bpoint = 0;
	double previous_popsize1 = popsizes1->at(0);
	double previous_popsize2 = popsizes2->at(0);
	int nlin1 = n1;
	int nlin2 = n2;

	//while (i < coals12.size()-2){
	////	int whichpop = coals12.at(i).second;
	//	if (whichpop == 1) nlin1 --;
	//	else if (whichpop == 2) nlin2--;
	//	i++;
	//}
	//i = coals12.size()-2;
	while (i < coals12.size()){
		bool print = false;
		if (i < 20) print = true;
		//
		// lk(i) = \frac{n1 choose 2}{N1_i} + \frac{n2 choose 2}{N2_i}  exp { - \int_{t_{i-2}}^t_i \frac{nlin1 choose 2}{N1_i} + frac{nlin2 choose 2}{N2_i} dt
		//
		if (print){
			cout << i << " line1 "<< nlin1 << " "<< nlin2 << " "<< previous_coaltime << " "<< previous_bpoint << " "<< previous_popsize1 << " "<< previous_popsize2 << "\n";
		}
		double current_coaltime = coals12.at(i).first;
		double current_bpoint = bpoints->at(j);
		double current_popsize1 = popsizes1->at(k);
		double current_popsize2 = popsizes2->at(k);
		if (print){
			cout << " current "<< current_coaltime << " n1:"<< nlin1 << " n2:"<< nlin2 << " "<< coals12[i].second << "\n";
			cout << " line2 " << nlin1<< " "<< current_coaltime << " "<< current_bpoint << " "<< current_popsize1 << " "<< current_popsize2 << "\n";
		}
		double add = 0;
		//cout << nlin<< " "<< st_nlin << " "<< coaltimes_in.size() << " "<< i <<"\n";
		//tree->print("test");
		double C1 = 0;
		double C2 = 0;
		if (nlin1 > 1) C1 = gsl_sf_fact(nlin1) / (gsl_sf_fact(2)* gsl_sf_fact(nlin1-2));
		if (nlin2 > 1) C2 = gsl_sf_fact(nlin2) / (gsl_sf_fact(2)* gsl_sf_fact(nlin2-2));
		//double C = gsl_sf_fact(nlin) / (gsl_sf_fact(2)* gsl_sf_fact(nlin-2));
		//cout << i << " C "<< C <<" nlin: "<< nlin << "\n";
		while (current_bpoint < current_coaltime){

			//add the bit after the previous event
			if (previous_coaltime < current_bpoint && previous_coaltime > previous_bpoint)	{
					//cout << "here1 " <<  current_bpoint << " "<< previous_coaltime << " "<< C1 << " "<< current_popsize1 << " "<< C2 << " "<< current_popsize2 << " \n";
					add += - (current_bpoint - previous_coaltime) * C1 / current_popsize1;
					add +=  - (current_bpoint - previous_coaltime) * C2 / current_popsize2;
					if (print) cout << " here1 "<< add << "\n";

			}
			else	{
				add += - (current_bpoint - previous_bpoint)* C1 / current_popsize1;
				add += - (current_bpoint - previous_bpoint)* C2 / current_popsize2;
				if (print) cout << " here2 "<< add << "\n";
				//cout << add<< " here2 " <<  current_bpoint << " "<< previous_coaltime << " "<< C1 << " "<< current_popsize1 << " "<< C2 << " "<< current_popsize2 << " \n";
			}
			//increment
			j++;
			k++;
			previous_bpoint = current_bpoint;
			previous_popsize1 = current_popsize1;
			previous_popsize2 = current_popsize2;
			current_bpoint = bpoints->at(j);
			current_popsize1 = popsizes1->at(k);
			current_popsize2 = popsizes2->at(k);

		}
		if (previous_coaltime > previous_bpoint){
			//cout << "here4 "<< current_coaltime << " "<< previous_coaltime << " "<< C << " "<< current_popsize << " \n";
			//cout << i  << " adding5 " <<-(current_coaltime - previous_coaltime)*  C/ current_popsize << "\n";
			add += -(current_coaltime - previous_coaltime)*  C1/ current_popsize1;
			add += -(current_coaltime - previous_coaltime)*  C2/ current_popsize2;
			if (print) cout << " here3 "<< add << "\n";
			//cout << add<< " here3 " <<  current_coaltime << " "<< previous_coaltime << " "<< C1 << " "<< current_popsize1 << " "<< C2 << " "<< current_popsize2 << " \n";
		}
		else {
			//cout << i << " adding6 " << - (current_coaltime - previous_bpoint) * C /current_popsize << "\n";
			add += - (current_coaltime - previous_bpoint) * C1 /current_popsize1;
			add += - (current_coaltime - previous_bpoint) * C2 /current_popsize2;
			if (print) cout << " here4 "<< add << "\n";
			//cout << add<< " here4 " <<  current_coaltime << " "<< previous_bpoint << " "<< C1 << " "<< current_popsize1 << " "<< C2 << " "<< current_popsize2 << " \n";
		}
		//cout << i << " adding7 " << log(C) - log(current_popsize) << "\n";
		double denom = C1/current_popsize1 + C2/current_popsize2;
		add += log(denom);
		if (print) cout << " here5 "<< add << "\n";
		//cout << " add1 "<< add << "\n";
		int whichpop = coals12.at(i).second;

		if (whichpop == 1){
			nlin1 --;
			add+= log(C1) - log(current_popsize1) - log(denom);
		}
		else {
			nlin2 --;
			add+= log(C2) - log(current_popsize2) - log(denom);
		}

		if (print){
			cout << " add "<<i << " "<< add << " "<< toreturn+add << "\n";
		}
		toreturn += add;
		//cout << i << " "<< add << "\n";
		previous_coaltime = current_coaltime; //increment, save coalescent time
		i = i+1;
	}
	// now all modeled coalescents have happened
	// if there is still more than one lineage, need to add likelihood of *no* additional coalescences in the remaining time
	//
	//cout << "toreturn? " << toreturn << "\n";
	//cout << "here\n"; cout.flush();
	//cout << st_nlin << " "<< coaltimes
	//tree->print("test");
	if (nlin1 > 1 or nlin2 > 1){
		//cout << "in here\n"; cout.flush();
		//cout << st_nlin << " "<< coaltimes_in.size() << "\n";
		double add = 0;
		double lastcoal;
		if (coals12.size() ==0 ) lastcoal = 0;
		else lastcoal= coals12.at(coals12.size()-1).first;
		double C1 = 0;
		double C2 = 0;
		if (nlin1 > 1) C1 = gsl_sf_fact(nlin1) / (gsl_sf_fact(2)* gsl_sf_fact(nlin1-2));
		if (nlin2 > 1) C2 = gsl_sf_fact(nlin2) / (gsl_sf_fact(2)* gsl_sf_fact(nlin2-2));

		for (int i = 0; i < bpoints->size(); i++){
			double dt;
			double N1 = popsizes1->at(i);
			double N2 = popsizes2->at(i);
			if (bpoints->at(i) < lastcoal) continue;
			if (i > 0 && bpoints->at(i) > lastcoal && bpoints->at(i-1) < lastcoal){
				dt = bpoints->at(i) - lastcoal;
			}
			else if (i > 0){
				dt = bpoints->at(i) - bpoints->at(i-1);
			}
			add +=  - dt *  ( C1/ N1 + C2/ N2);
		}
		toreturn += add;
	}
	///cout << "not here\n"; cout.flush();
	return toreturn;

}
void TreeState::print(){
	string toprint = tree->get_newick_format();
	cout << toprint << "\n";
}

double TreeState::structured_llk_binom(double p){
	double toreturn;
	int total = n1+n2;
	toreturn+= gsl_sf_lnfact(total);
	toreturn -=gsl_sf_lnfact(n1);
	toreturn -=gsl_sf_lnfact(n2);
	toreturn += n1 *log(p);
	toreturn += n2 *log(1-p);
	return toreturn;
}

double TreeState::getp(){
	double toreturn;
	toreturn = (double) n1 / ((double) n1+ (double)n2);
	return toreturn;
}


void TreeState::set_UPGMA(){
	tree->set_UPGMA(aln);
}
